import torch.nn as nn
import torch

from .BackBone import (
    VGG_Backbone,
    ResNet_Backbone,
    DenseNet_Backbone,
    EfficientNet_Backbone,
    ViT_Backbone,
    SwinTransformer_Backbone,
    UNet_Backbone,

)

import torch
import torch.nn as nn
import torch.nn.functional as F

class Siamese(nn.Module):
    def __init__(self, batch_size=64, channel=3, input_shape=[64, 3, 105, 105], 
                 backbone='efficientnet', backbone_version='efficientnet_b0', 
                 pretrained=False, distance_type='cosine'):
        super(Siamese, self).__init__()
        
        # 选择骨干网络
        if backbone == 'vgg':
            self.backbone = VGG_Backbone(backbone_version, pretrained)
        elif backbone == 'resnet':
            self.backbone = ResNet_Backbone(backbone_version, pretrained)
        elif backbone == 'densenet':
            self.backbone = DenseNet_Backbone(backbone_version, pretrained)
        elif backbone == 'efficientnet':
            self.backbone = EfficientNet_Backbone(backbone_version, pretrained)
        elif backbone == 'vit':
            self.backbone = ViT_Backbone(backbone_version, pretrained)
        elif backbone == 'swin':
            self.backbone = SwinTransformer_Backbone(backbone_version, pretrained)
        elif backbone == 'unet':
            self.backbone = UNet_Backbone()
        else:
            raise ValueError(f"Unsupported backbone: {backbone}")
        
        # 计算特征图展平后的维度
        self.flat_shape = self._get_flat_shape(batch_size=batch_size, channel=channel, input_shape=input_shape)

        # 定义全连接层
        self.fully_connect1 = nn.Linear(1, 512)  # 输入是距离（标量）
        self.bn1 = nn.BatchNorm1d(512)
        self.dropout1 = nn.Dropout(0.5)
        self.fully_connect2 = nn.Linear(512, 1)
        self.relu = nn.ReLU()

        # 距离类型
        self.distance_type = distance_type

        # 初始化协方差矩阵的逆矩阵（仅用于马氏距离）
        if distance_type == 'mahalanobis':
            self.register_buffer('cov_inv', torch.eye(self.flat_shape))  # 初始化为单位矩阵

    def _get_flat_shape(self, batch_size, channel, input_shape):
        # 使用一个输入来计算特征图展平后的维度
        dummy_input = torch.randn(batch_size, channel, input_shape[0], input_shape[1])
        features = self.backbone(dummy_input)
        features = torch.flatten(features, 1)
        return features.view(features.size(0), -1).size(1)

    def update_cov_inv(self, features):
        """
        更新协方差矩阵的逆矩阵（仅用于马氏距离）
        :param features: 所有训练数据的特征向量 (N, D)，N是样本数，D是特征维度
        """
        if self.distance_type != 'mahalanobis':
            raise ValueError("update_cov_inv is only applicable for Mahalanobis distance.")
        
        cov_matrix = torch.cov(features.T)  # 计算协方差矩阵 (D, D)
        self.cov_inv = torch.inverse(cov_matrix + 1e-6 * torch.eye(cov_matrix.size(0)))  # 添加小量防止奇异矩阵

    def _compute_distance(self, x1, x2):
        """
        根据 distance_type 计算距离
        :param x1: 特征向量1 (batch_size, D)
        :param x2: 特征向量2 (batch_size, D)
        :return: 距离 (batch_size, 1)
        """
        if self.distance_type == 'l1':
            # L1距离
            return torch.abs(x1 - x2).sum(dim=1, keepdim=True)
        elif self.distance_type == 'euclidean':
            # 欧氏距离
            return torch.norm(x1 - x2, p=2, dim=1, keepdim=True)
        elif self.distance_type == 'mahalanobis':
            # 马氏距离
            diff = x1 - x2
            dist = torch.sqrt(torch.sum((diff @ self.cov_inv) * diff, dim=1, keepdim=True))
            return dist
        elif self.distance_type == 'cosine':
            # 余弦相似度（转换为距离：1 - 余弦相似度）
            cosine_sim = F.cosine_similarity(x1, x2, dim=1).unsqueeze(1)
            return 1 - cosine_sim
        else:
            raise ValueError(f"Unsupported distance type: {self.distance_type}")

    def forward(self, x):
        x1 = x[0]
        x2 = x[1]

        # 通过骨干网络提取特征
        x1 = self.backbone(x1)
        x2 = self.backbone(x2)
        
        # 展平特征图
        x1 = torch.flatten(x1, 1)
        x2 = torch.flatten(x2, 1)
        
        # 计算距离
        x = self._compute_distance(x1, x2)
        
        # 通过全连接层
        x = self.fully_connect1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.dropout1(x)
        
        x = self.fully_connect2(x)
        return x



if __name__ =="__main__":
    # 假设输入图像的形状为 (3, 224, 224)
    input_shape=[10, 3, 224, 224]

    # 使用DenseNet121作为骨干网络
    model_densenet = Siamese(input_shape=input_shape, backbone='efficientnet', backbone_version='efficientnet_b0', pretrained=False)

    # 假设输入是两个大小为 (batch_size, 3, 224, 224) 的图像
    x1 = torch.randn(10, 3, 224, 224)  # batch_size=10
    x2 = torch.randn(10, 3, 224, 224)

    # 使用DenseNet121模型
    output_densenet = model_densenet((x1, x2))
    print(output_densenet.shape)  # 输出形状应为 (batch_size, 1)


    # backbone_version='efficientnet_b0'

    # backbone = EfficientNet_Backbone(backbone_version)
    # input_shape=[224,224]
    # # flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
    # # print(flat_shape)

    #     # 使用一个虚拟输入来计算特征图展平后的维度
    # dummy_input = torch.randn(10,3, *input_shape)
    # features = backbone(dummy_input)
    
    # # 添加全局平均池化层，将特征图统一到 (batch_size, channels, 1, 1)
    # # features = nn.AdaptiveAvgPool2d((1, 1))(features)
    # features = torch.flatten(features, 1)
    # print(features.view(features.size(0), -1).size(1))

